use std::alloc::{alloc, dealloc, Layout};
use std::ptr;

const MALLOC_PATTERN: u8 = 0xAA;
const FREE_PATTERN: u8 = 0xBB;
const ALLOC_TEST_MAGIC: u32 = 0x12345678;

struct BlockHeader {
    magic_number: u32,
    bytes: usize,
}

static mut ALLOCATED_BYTES: usize = 0;
static mut ALLOCATION_LIMIT: isize = -1;

fn alloc_test_get_header(ptr: *mut u8) -> *mut BlockHeader {
    unsafe {
        let header_ptr = ptr.offset(-std::mem::size_of::<BlockHeader>() as isize);
        assert!((*header_ptr).magic_number == ALLOC_TEST_MAGIC);
        header_ptr as *mut BlockHeader
    }
}

fn alloc_test_overwrite(ptr: *mut u8, length: usize, pattern: u8) {
    unsafe {
        ptr.write_bytes(pattern, length);
    }
}

fn alloc_test_malloc(bytes: usize) -> *mut u8 {
    if unsafe { ALLOCATION_LIMIT } == 0 {
        return ptr::null_mut();
    }

    let layout = Layout::from_size_align_unchecked(
        std::mem::size_of::<BlockHeader>() + bytes,
        std::mem::align_of::<BlockHeader>(),
    );

    let header_ptr = unsafe { alloc(layout) };
    if header_ptr.is_null() {
        return ptr::null_mut();
    }

    unsafe {
        (*header_ptr).magic_number = ALLOC_TEST_MAGIC;
        (*header_ptr).bytes = bytes;

        let ptr = header_ptr.offset(std::mem::size_of::<BlockHeader>() as isize);
        alloc_test_overwrite(ptr, bytes, MALLOC_PATTERN);

        ALLOCATED_BYTES += bytes;

        if ALLOCATION_LIMIT > 0 {
            ALLOCATION_LIMIT -= 1;
        }

        ptr as *mut u8
    }
}

fn alloc_test_free(ptr: *mut u8) {
    if ptr.is_null() {
        return;
    }

    let header_ptr = unsafe { alloc_test_get_header(ptr) };
    let block_size = unsafe { (*header_ptr).bytes };

    assert!(unsafe { ALLOCATED_BYTES } >= block_size);

    unsafe {
        alloc_test_overwrite(ptr, block_size, FREE_PATTERN);
        (*header_ptr).magic_number = 0;
        dealloc(header_ptr as *mut u8, Layout::from_size_align_unchecked(
            std::mem::size_of::<BlockHeader>() + block_size,
            std::mem::align_of::<BlockHeader>(),
        ));
    }

    unsafe {
        ALLOCATED_BYTES -= block_size;
    }
}

fn alloc_test_realloc(ptr: *mut u8, bytes: usize) -> *mut u8 {
    let new_ptr = alloc_test_malloc(bytes);
    if new_ptr.is_null() {
        return ptr::null_mut();
    }

    if !ptr.is_null() {
        let header_ptr = unsafe { alloc_test_get_header(ptr) };
        let bytes_to_copy = unsafe { (*header_ptr).bytes };

        let bytes_to_copy = if bytes_to_copy > bytes {
            bytes
        } else {
            bytes_to_copy
        };

        unsafe {
            ptr::copy_nonoverlapping(ptr, new_ptr, bytes_to_copy);
            alloc_test_free(ptr);
        }
    }

    new_ptr
}

fn alloc_test_calloc(nmemb: usize, bytes: usize) -> *mut u8 {
    let total_bytes = nmemb * bytes;
    let result = alloc_test_malloc(total_bytes);

    if !result.is_null() {
        unsafe {
            ptr::write_bytes(result, 0, total_bytes);
        }
    }

    result
}

fn alloc_test_strdup(string: &str) -> *mut u8 {
    let result = alloc_test_malloc(string.len() + 1);

    if !result.is_null() {
        unsafe {
            ptr::copy_nonoverlapping(string.as_ptr(), result, string.len());
            *(result.offset(string.len() as isize)) = 0;
        }
    }

    result
}

fn alloc_test_set_limit(alloc_count: isize) {
    unsafe {
        ALLOCATION_LIMIT = alloc_count;
    }
}

fn alloc_test_get_allocated() -> usize {
    unsafe { ALLOCATED_BYTES }
}
